from scipy.stats import mode
from sklearn.datasets import load_iris
from sklearn.mixture import GaussianMixture
import matplotlib.pyplot as plt

iris = load_iris()
model = GaussianMixture(n_components=3)
predict = model.fit_predict(iris.data)


def score(predict, gt):
    assert len(predict) == len(gt)
    map_ = {}
    for c in set(predict):
        map_[c] = mode(gt[predict == c], keepdims=True)[0][0]
    correct_count = sum(1 for i in range(len(predict)) if map_[predict[i]] == gt[i])
    score_value = correct_count / len(predict)
    return score_value


def ke_shi_hua():
    fig, axs = plt.subplots(1, 2, figsize=(13, 7))
    axs[0].set_title("ground truth")
    axs[1].set_title("prediction")
    for target in range(3):
        axs[0].scatter(iris.data[iris.target == target, 1], iris.data[iris.target == target, 3])
        axs[1].scatter(iris.data[predict == target, 1], iris.data[predict == target, 3])
    plt.show()


if __name__ == '__main__':
    print(f'GMM模型的purity评分为：{score(predict, iris.target)}')
    ke_shi_hua()
